import os
import json
import random
from glob import glob

"""
生成文本检测训练数据

"""

# ================== 配置参数 ==================
json_root_folder = '/root/PaddleOCR/dataset/en_det_300_res'      # 替换为你的根文件夹路径
output_dir = '/root/PaddleOCR/dataset/en_det_300_res'          # 输出文件夹
image_prefix = '1'                            # 图片路径前缀，如 images/

# 创建输出目录
os.makedirs(output_dir, exist_ok=True)

# 输出文件路径
train_txt = os.path.join(output_dir, 'train.txt')
val_txt = os.path.join(output_dir, 'val.txt')

# ================== 1. 递归查找所有 .json 文件 ==================
print("正在搜索所有 JSON 文件...")
json_files = []
for root, dirs, files in os.walk(json_root_folder):
    for file in files:
        if file.lower().endswith('.json'):
            json_files.append(os.path.join(root, file))

print(f"共找到 {len(json_files)} 个 JSON 文件。")

# ================== 2. 打乱顺序并划分 9:1 ==================
random.seed(42)  # 固定随机种子，保证可复现
random.shuffle(json_files)

split_idx = int(0.9 * len(json_files))
train_files = json_files[:split_idx]
val_files = json_files[split_idx:]

print(f"训练集: {len(train_files)} 个文件")
print(f"验证集: {len(val_files)} 个文件")

# ================== 3. 处理函数 ==================
def process_json_file(json_path):
    try:
        with open(json_path, 'r', encoding='utf-8') as f:
            data = json.load(f)
    except Exception as e:
        print(f"读取失败: {json_path}, 错误: {e}")
        return None

    # 获取 imagePath，若不存在则用文件名
    image_filename = data.get("imagePath")
    if not image_filename:
        # 使用 JSON 文件同名的图片
        base_name = os.path.splitext(os.path.basename(json_path))[0]
        image_filename = f"{base_name}.png"  # 可根据实际图片格式调整

    # image_path = f"{image_prefix}/{image_filename}"
    image_path = json_path.split(json_root_folder)[-1][1:].replace(".json",".png")
    print(image_path)

    # 提取 shapes
    entries = []
    for shape in data.get("shapes", []):
        desc = str(shape.get("description", "")).strip()  # 转字符串并去空
        transcription = desc if desc else "###"

        # 坐标取整
        points = [[int(round(float(coord))) for coord in point] for point in shape["points"]]

        entries.append({
            "transcription": transcription,
            "points": points
        })

    # 格式化为一行字符串
    line = image_path + "\t" + json.dumps(entries, ensure_ascii=False, separators=(',', ':'))
    return line

# ================== 4. 写入输出文件 ==================
def write_lines(file_list, output_path):
    with open(output_path, 'w', encoding='utf-8') as f:
        for json_path in file_list:
            line = process_json_file(json_path)
            if line:
                f.write(line + '\n')
    print(f"已写入: {output_path},数量：{len(file_list)}")

# 处理并写入
write_lines(train_files, train_txt)
write_lines(val_files, val_txt)

print("✅ 所有文件处理完成！")